import numpy as np
from datetime import date, timedelta
from xgboost import XGBRegressor
from cachetools import TTLCache

prediction_cache = TTLCache(maxsize=1000000, ttl=86400)   # 历史数据缓存1小时

def train_and_predict(ticker):
    print('开始批量训练和预测')
    cache_key = f"predictions_{ticker['股票代码']}"  # 用股票代码作为缓存键
    cached_result = prediction_cache.get(cache_key)

    if cached_result is not None:
        print(f"从缓存中获取股票 {ticker['股票代码']} 的预测结果")
        return cached_result

    # 获取股票数据
    stock_data = ticker

    if stock_data.empty:
        print("没有获取到有效的股票数据")
        return []

    # 按照股票代码分组处理
    grouped = stock_data.groupby('股票代码')
    all_predictions = []
    today = date.today()

    for ticker, group in grouped:
        try:
            # 确保数据足够
            if len(group) < 10:
                print(f"股票 {ticker} 数据不足，跳过")
                continue

            # 准备数据
            df = group.copy()
            df['day'] = range(len(df))

            X = df[['day']].values
            y = df['收盘价'].values
            y_v = np.var(y)

            # 训练模型
            model = XGBRegressor(n_estimators=50)
            model.fit(X, y)

            # 预测未来30天
            future_days = np.array(range(len(df), len(df) + 10)).reshape(-1, 1)
            preds = model.predict(future_days)

            # 生成预测结果
            predictions = [{
                'ticker': ticker,
                'date': today + timedelta(days=i),
                'predicted_close': round(float(preds[i])+np.random.normal(0, np.sqrt(y_v)*0.3),6)
            } for i in range(10)]

            all_predictions.extend(predictions)
            #print(predictions)
            #print(f"股票 {ticker} 预测完成")

        except Exception as e:
            #print(f"股票 {ticker} 预测失败: {str(e)}")
            continue
    #print(all_predictions)
    prediction_cache[cache_key] = all_predictions
    print(f"完成所有预测，共 {len(all_predictions)} 条预测结果")
    return all_predictions